import os
import json
import requests
from typing import Dict, List, Any, Optional, Union
from dotenv import load_dotenv
from loguru import logger
from datetime import datetime, date

# 加载环境变量
load_dotenv()

class DeepSeekClient:
    """DeepSeek大模型客户端，用于自然语言处理和SQL生成"""
    
    def __init__(self):
        self.api_key = os.getenv("DEEPSEEK_API_KEY")
        self.model = os.getenv("DEEPSEEK_MODEL", "deepseek-chat")
        self.embedding_model = os.getenv("DEEPSEEK_EMBEDDING_MODEL", "text-embedding")
        self.api_base = os.getenv("DEEPSEEK_API_BASE", "https://api.deepseek.com/v1")
        
        logger.info(f"初始化DeepSeek客户端: model={self.model}, api_base={self.api_base}")
        
        if not self.api_key:
            logger.warning("未配置DeepSeek API密钥，某些功能可能不可用")
    
    def _make_api_request(self, endpoint: str, data: Dict[str, Any]) -> Dict[str, Any]:
        """向DeepSeek API发送请求"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        url = f"{self.api_base}/{endpoint}"
        
        # 打印请求信息（隐藏API密钥）
        safe_headers = headers.copy()
        safe_headers["Authorization"] = "Bearer sk-***" 
        logger.debug(f"API请求: URL={url}, Headers={safe_headers}")
        logger.debug(f"请求数据: {json.dumps(data, ensure_ascii=False)[:1000]}...")
        
        response = None
        try:
            response = requests.post(url, headers=headers, json=data)
            
            # 打印响应状态和头信息
            logger.debug(f"API响应状态: {response.status_code}")
            logger.debug(f"API响应头: {response.headers}")
            
            # 尝试获取响应内容
            response_text = response.text
            logger.debug(f"API原始响应内容: {response_text[:2000]}...")
            
            # 确保响应是有效的JSON格式
            try:
                response_json = response.json()
                logger.info(f"API请求成功: {endpoint}")
                return response_json
            except json.JSONDecodeError as e:
                logger.error(f"API响应不是有效的JSON格式: {str(e)}")
                logger.error(f"响应内容: {response_text}")
                raise ValueError(f"API响应不是有效的JSON格式: {str(e)}")
                
        except requests.exceptions.RequestException as e:
            logger.error(f"API请求失败: {str(e)}")
            if response:
                logger.error(f"API响应状态码: {response.status_code}")
                logger.error(f"API响应内容: {response.text if hasattr(response, 'text') else '无响应内容'}")
            raise
    
    def generate_sql(self, natural_language_query: str, database_schema: Dict[str, Any], 
                    example_data: Optional[Dict[str, List[Dict[str, Any]]]] = None) -> Dict[str, Any]:
        """根据自然语言查询生成SQL"""
        logger.info(f"生成SQL查询: '{natural_language_query}'")
        
        # 构建提示信息 - 修改为更简洁的指令
        system_prompt = """你是一个专业的SQL生成助手。你的任务是根据用户的自然语言查询和数据库架构，生成精确的SQL查询语句。
        请直接返回SQL查询语句，不要添加任何解释、注释或额外内容。只输出有效的SQL语句。"""
        
        # 创建数据库架构的文本描述
        schema_text = "数据库架构:\n"
        for table_name, table_info in database_schema.items():
            schema_text += f"表名: {table_name}\n"
            schema_text += "列:\n"
            for column in table_info["schema"]:
                schema_text += f"  - {column['name']} ({column['type']})"
                if not column["nullable"]:
                    schema_text += " NOT NULL"
                if column.get("default"):
                    schema_text += f" DEFAULT {column['default']}"
                schema_text += "\n"
            
            if table_info["relationships"]:
                schema_text += "外键关系:\n"
                for rel in table_info["relationships"]:
                    schema_text += f"  - {','.join(rel['constrained_columns'])} -> {rel['referred_table']}.{','.join(rel['referred_columns'])}\n"
            schema_text += "\n"
        
        # 如果有示例数据，添加到提示中
        example_data_text = ""
        if example_data:
            example_data_text = "示例数据:\n"
            for table_name, rows in example_data.items():
                example_data_text += f"表 {table_name} 的示例数据:\n"
                if rows:
                    for i, row in enumerate(rows[:3]):  # 只显示前3行
                        example_data_text += f"  行 {i+1}: {json.dumps(row, ensure_ascii=False)}\n"
                example_data_text += "\n"
        
        # 用户查询
        user_query = f"自然语言查询: {natural_language_query}\n\n请生成对应的SQL查询:"
        
        # 构建完整消息
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"{schema_text}\n{example_data_text}\n{user_query}"}
        ]
        
        # 调用API时添加额外参数，避免序列化问题
        try:
            # 打印模型名称和完整请求参数
            logger.info(f"使用模型: {self.model}")
            
            response = self._make_api_request("chat/completions", {
                "model": self.model,
                "messages": messages,
                "temperature": 0.1,  # 较低的温度以获得更确定性的输出
                "max_tokens": 1000   # 添加max_tokens参数
            })
            
            # 提取生成的SQL
            if "choices" not in response or not response["choices"]:
                logger.error(f"API响应中缺少choices字段: {response}")
                return {
                    "original_query": natural_language_query,
                    "generated_sql": "",
                    "error": "API响应格式不正确，缺少choices字段",
                    "raw_response": response
                }
                
            message = response["choices"][0].get("message", {})
            if not message or "content" not in message:
                logger.error(f"API响应中缺少message.content字段: {response}")
                return {
                    "original_query": natural_language_query,
                    "generated_sql": "",
                    "error": "API响应格式不正确，缺少message.content字段",
                    "raw_response": response
                }
                
            generated_sql = message.get("content", "")
            
            # 清理SQL（移除代码块标记等）
            sql = self._clean_sql(generated_sql)
            logger.info(f"生成的SQL: {sql}")
            
            return {
                "original_query": natural_language_query,
                "generated_sql": sql,
                "raw_response": response
            }
        except Exception as e:
            logger.error(f"生成SQL失败: {str(e)}")
            return {
                "original_query": natural_language_query,
                "generated_sql": "",
                "error": str(e)
            }
    
    def _clean_sql(self, sql_text: str) -> str:
        """清理从模型返回的SQL文本"""
        # 移除Markdown代码块标记
        sql_text = sql_text.replace("```sql", "").replace("```", "")
        
        # 移除前后多余的空白字符
        sql_text = sql_text.strip()
        
        return sql_text
    
    def explain_sql(self, sql_query: str) -> str:
        """解释SQL查询的含义"""
        logger.info(f"解释SQL查询: '{sql_query}'")
        
        # 修改为更简洁的指令
        system_prompt = """简要解释SQL查询的作用。使用直接、简洁的语言。不要使用比喻，不要过度解释概念，不要进行教学。
        请用1-2句话总结查询的目的，不需要解释SQL语法细节。回答应该简短精确，直接说明这个查询会返回什么结果。"""
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"请简洁解释以下SQL查询:\n\n{sql_query}"}
        ]
        
        try:
            response = self._make_api_request("chat/completions", {
                "model": self.model,
                "messages": messages,
                "max_tokens": 100   # 减少max_tokens以鼓励简洁回答
            })
            
            if "choices" not in response or not response["choices"]:
                logger.error(f"API响应中缺少choices字段: {response}")
                return "无法解释SQL: API响应格式不正确"
                
            explanation = response.get("choices", [{}])[0].get("message", {}).get("content", "")
            logger.info(f"SQL解释已生成: {explanation[:100]}...")
            return explanation
        except Exception as e:
            logger.error(f"解释SQL失败: {str(e)}")
            return f"无法解释SQL: {str(e)}"
    
    def analyze_query_results(self, query_result: List[Dict[str, Any]], 
                             original_query: str) -> Dict[str, Any]:
        """分析查询结果并提供见解"""
        logger.info(f"分析查询结果: '{original_query}'")
        
        # 确保query_result是有效的列表
        if not isinstance(query_result, list):
            logger.error(f"查询结果必须是列表，收到的类型: {type(query_result)}")
            return {"analysis": "无法分析结果：数据格式不正确"}
            
        # 限制传递的数据量
        safe_result = query_result[:10] if query_result else []
        result_json = json.dumps(safe_result, ensure_ascii=False)
        logger.debug(f"查询结果示例: {result_json[:500]}...")
        
        # 修改为更简洁的指令
        system_prompt = """简洁分析查询结果，直接给出关键数字和发现。限制在2-3句话内，只突出最重要的事实和数字。
        不要提供背景信息、解释原理、建议或多余分析。只提供结果中直接可见的关键信息，用最简单直接的语言。"""
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"查询: {original_query}\n\n查询结果: {result_json}\n\n请简洁分析:"}
        ]
        
        try:
            response = self._make_api_request("chat/completions", {
                "model": self.model,
                "messages": messages,
                "max_tokens": 100   # 减少max_tokens以鼓励简洁回答
            })
            
            if "choices" not in response or not response["choices"]:
                logger.error(f"API响应中缺少choices字段: {response}")
                return {
                    "original_query": original_query,
                    "analysis": "无法分析结果: API响应格式不正确",
                    "raw_response": response
                }
                
            analysis = response.get("choices", [{}])[0].get("message", {}).get("content", "")
            logger.info(f"分析结果已生成: {analysis[:100]}...")
            
            return {
                "original_query": original_query,
                "analysis": analysis
            }
        except Exception as e:
            logger.error(f"分析查询结果失败: {str(e)}")
            return {
                "original_query": original_query,
                "analysis": f"无法分析结果: {str(e)}"
            }
    
    def suggest_visualization(self, query_result: List[Dict[str, Any]], 
                             query: str) -> Dict[str, Any]:
        """基于查询结果推荐可视化方式"""
        logger.info(f"推荐可视化方式: '{query}'")
        
        # 确保query_result是有效的列表
        if not isinstance(query_result, list):
            logger.error(f"查询结果必须是列表，收到的类型: {type(query_result)}")
            return {"visualization_suggestion": "无法推荐可视化：数据格式不正确"}
            
        # 限制传递的数据量
        safe_result = query_result[:10] if query_result else []
        result_json = json.dumps(safe_result, ensure_ascii=False)
        logger.debug(f"查询结果示例: {result_json[:500]}...")
        
        # 修改为更简洁的指令
        system_prompt = """用一句话推荐最适合的可视化图表类型。
        只需直接说明图表类型（如柱状图、折线图、饼图等），没有其他内容。
        如果是单值结果，只需说明显示数字即可，不需要详细解释或建议。"""
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"查询: {query}\n\n查询结果: {result_json}\n\n一句话推荐可视化方式:"}
        ]
        
        try:
            response = self._make_api_request("chat/completions", {
                "model": self.model,
                "messages": messages,
                "max_tokens": 50   # 减少max_tokens以鼓励简洁回答
            })
            
            if "choices" not in response or not response["choices"]:
                logger.error(f"API响应中缺少choices字段: {response}")
                return {
                    "query": query,
                    "visualization_suggestion": "无法推荐可视化: API响应格式不正确",
                    "raw_response": response
                }
                
            suggestion = response.get("choices", [{}])[0].get("message", {}).get("content", "")
            logger.info(f"可视化建议已生成: {suggestion[:100]}...")
            
            return {
                "query": query,
                "visualization_suggestion": suggestion
            }
        except Exception as e:
            logger.error(f"推荐可视化方式失败: {str(e)}")
            return {
                "query": query,
                "visualization_suggestion": f"无法推荐可视化: {str(e)}"
            }
    
    def get_embedding(self, text: str) -> List[float]:
        """获取文本的嵌入向量表示"""
        logger.info(f"获取文本嵌入向量: '{text[:50]}...'")
        
        if not self.api_key:
            logger.error("无法生成嵌入：未配置API密钥")
            return []
        
        try:
            response = self._make_api_request("embeddings", {
                "model": self.embedding_model,
                "input": text,
                "encoding_format": "float"  # 确保返回格式正确
            })
            
            if "data" not in response or not response["data"]:
                logger.error(f"API响应中缺少data字段: {response}")
                return []
                
            embedding = response.get("data", [{}])[0].get("embedding", [])
            logger.info(f"成功获取嵌入向量，维度: {len(embedding)}")
            return embedding
        except Exception as e:
            logger.error(f"生成嵌入失败: {str(e)}")
            return []


# 单例模式
_model_client = None

def get_model_client() -> DeepSeekClient:
    """获取DeepSeek模型客户端的单例实例"""
    global _model_client
    if _model_client is None:
        _model_client = DeepSeekClient()
    return _model_client 